"""
1.	Darknet-19。YOLOv2使用了一个新的分类网络作为特征提取部分，参考了前人的工作经验。类似于VGG，网络使用了较多的3 * 3卷积核，
在每一次池化操作后把通道数翻倍。借鉴了network in network的思想，网络使用了全局平均池化（global average pooling）做预测，
把1 * 1的卷积核置于3 * 3的卷积核之间，用来压缩特征。使用batch normalization稳定模型训练，加速收敛，正则化模型。
最终得出的基础模型就是Darknet-19，包含19个卷积层、5个最大值池化层（max pooling layers ）。按下面要求实现代码：
"""
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers, activations, losses, optimizers, metrics

# ①	定义主程序入口
if '__main__' == __name__:

    # ②	定义类函数Conv_BN_LeakyReLU()完成卷积、BN层、激活
    def Conv_BN_LeakyReLU(filters, ksize=(3, 3), strides=(1, 1)):
        return keras.Sequential([
            layers.Conv2D(filters, ksize, strides, 'same'),
            layers.BatchNormalization(),
            layers.ReLU(),
        ])

    # ③	定义类函数DarkNet_19()，参考下图
    class DarkNet_19(keras.Model):

        # ④	用nn.Sequential()，定义6个卷积块
        def __init__(self, **kwargs):
            super().__init__(**kwargs)
            self.layer01 = keras.Sequential([
                Conv_BN_LeakyReLU(32),
                layers.MaxPool2D((2, 2), (2, 2), 'same')
            ])
            self.layer02 = keras.Sequential([
                Conv_BN_LeakyReLU(64),
                layers.MaxPool2D((2, 2), (2, 2), 'same')
            ])
            self.layer03 = keras.Sequential([
                Conv_BN_LeakyReLU(128),
                Conv_BN_LeakyReLU(64, (1, 1)),
                Conv_BN_LeakyReLU(128),
                layers.MaxPool2D((2, 2), (2, 2), 'same')
            ])
            self.layer04 = keras.Sequential([
                Conv_BN_LeakyReLU(256),
                Conv_BN_LeakyReLU(128, (1, 1)),
                Conv_BN_LeakyReLU(256),
                layers.MaxPool2D((2, 2), (2, 2), 'same')
            ])
            self.layer05 = keras.Sequential([
                Conv_BN_LeakyReLU(512),
                Conv_BN_LeakyReLU(256, (1, 1)),
                Conv_BN_LeakyReLU(512),
                Conv_BN_LeakyReLU(256, (1, 1)),
                Conv_BN_LeakyReLU(512),
                layers.MaxPool2D((2, 2), (2, 2), 'same')
            ])
            self.layer06 = keras.Sequential([
                Conv_BN_LeakyReLU(1024),
                Conv_BN_LeakyReLU(512, (1, 1)),
                Conv_BN_LeakyReLU(1024),
                Conv_BN_LeakyReLU(512, (1, 1)),
                Conv_BN_LeakyReLU(1024),
            ])
            self.top = keras.Sequential([
                Conv_BN_LeakyReLU(1000, (1, 1)),
                layers.GlobalAvgPool2D(),
                layers.Softmax(),
            ])

        # ⑤	进行前向传播
        def call(self, x, training=None):
            x = self.layer01(x)
            print('x.shape', tf.shape(x))
            x = self.layer02(x)
            print('x.shape', tf.shape(x))
            x = self.layer03(x)
            print('x.shape', tf.shape(x))
            x = self.layer04(x)
            print('x.shape', tf.shape(x))
            x = self.layer05(x)
            print('x.shape', tf.shape(x))
            x = self.layer06(x)
            print('x.shape', tf.shape(x))
            x = self.top(x)
            return x

    model = DarkNet_19()
    model.build((None, 224, 224, 3))
    model.summary()

    # ⑥	自定义输入张量(1, 3, 224, 224)，验证网络可以正常跑通
    x = tf.zeros((1, 224, 224, 3), dtype=tf.float32)
    pred = model(x)

    # ⑦	在前向传播中，打印6个卷积块的输出维度
    # ⑧	打印网络最后的输出维度
    print('Final output:', pred.shape)
